package com.imooc.datasource;

import com.alibaba.druid.pool.DruidDataSource;
import com.imooc.constants.DataSourceConstant;
import com.imooc.provider.DynamicDataSourceProvider;
import com.imooc.util.DynamicDataSourceContextHolder;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.ReflectionUtils;

import javax.sql.DataSource;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @Author Tim
 * @Description:
 * @Version 1.0
 */
@Slf4j
public class DynamicRoutingDataSource extends AbstractRoutingDataSource implements InitializingBean, DisposableBean {

    /**
     * 加载数据源集合
     */
    private final List<DynamicDataSourceProvider> providers;

    /**
     * 默认数据源的名称----> master,允许用户进行修改
     */
    @Setter
    private String primary = DataSourceConstant.PRIMARY;

    /**
     * 所有的数据源
     */
    private final Map<String, DataSource> allDataSources = new ConcurrentHashMap<>();

    public DynamicRoutingDataSource(List<DynamicDataSourceProvider> providers) {
        this.providers = providers;
    }

    @Override
    protected DataSource determineDataSource() {
        String currentDataSource = DynamicDataSourceContextHolder.getCurrentDataSource();
        return getDataSource(currentDataSource);
    }

    private DataSource getDataSource(String dataSource) {
        if (StringUtils.isEmpty(dataSource)) {
            return getDefaultDataSource();
        } else if (allDataSources.containsKey(dataSource)) {
            return allDataSources.get(dataSource);
        }
        return getDefaultDataSource();
    }

    /**
     * 返回默认数据源
     *
     * @return 返回默认数据源
     */
    private DataSource getDefaultDataSource() {
        DataSource dataSource = allDataSources.get(primary);
        if (dataSource != null) {
            return dataSource;
        }
        throw new RuntimeException("No default data source is currently found");
    }


    public void addDataSource(String ds, DataSource dataSource) {
        //如果以前存在这个key。把就得这个key对应的value 返回.
        DataSource oldDataSource = allDataSources.put(ds, dataSource);
        if (oldDataSource != null) {
            closeDataSource(ds, oldDataSource);
        }
        log.info("Add a data source {} success", ds);
    }

    /**
     * 关闭数据源
     *
     * @param ds
     * @param dataSource
     */
    private void closeDataSource(String ds, DataSource dataSource) {
        if (dataSource instanceof DruidDataSource dds) {
            dds.close();
        } else {
            try {
                Method close = ReflectionUtils.findMethod(dataSource.getClass(), "close");
                if (Objects.nonNull(close)) {
                    close.invoke(dataSource);
                }
            } catch (Exception e) {
                log.error("关闭数据源 {} 失败, e=", ds, e);
            }
        }
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        Map<String, DataSource> dataSourceMap = new HashMap<>(16);
        providers.stream().map(DynamicDataSourceProvider::loadDataSources).forEach(dataSourceMap::putAll);
        for (Map.Entry<String, DataSource> entry : dataSourceMap.entrySet()) {
            addDataSource(entry.getKey(), entry.getValue());
        }
    }

    @Override
    public void destroy() throws Exception {
        log.info("close all dataSource");
        for (Map.Entry<String, DataSource> entry : allDataSources.entrySet()) {
            closeDataSource(entry.getKey(), entry.getValue());
        }
    }


}
